#!/usr/bin/env python3
"""
Test script for pixi signal handling.

Tests the following scenarios:
1. Non-interactive mode + kill -INT → should forward (child receives signal)
2. Non-interactive mode + kill -TERM → should forward (child receives signal)
3. Interactive mode + kill -TERM → should forward (SIGTERM always forwards)
4. Interactive mode + kill -INT + same PGID → should NOT forward (trade-off, like UV)

Run with: python test_signals.py
"""

import os
import signal
import subprocess
import sys
import tempfile
import time
from pathlib import Path

# Colors for output
GREEN = "\033[92m"
RED = "\033[91m"
YELLOW = "\033[93m"
RESET = "\033[0m"

# Get the pixi binary path
PIXI_BIN = Path(__file__).parent / "target" / "pixi" / "debug" / "pixi"
if not PIXI_BIN.exists():
    PIXI_BIN = Path(__file__).parent / "target" / "debug" / "pixi"
if not PIXI_BIN.exists():
    print(f"{RED}Error: Could not find pixi binary. Run 'cargo build' first.{RESET}")
    sys.exit(1)

# Test child script that catches signals
CHILD_SCRIPT = """
import os
import signal
import sys
import time
from pathlib import Path

# Output file is passed via environment variable
output_file = Path(os.environ.get("OUTPUT_FILE", "output.txt"))

def handler(signum, frame):
    print(f"Got signal {signum}", flush=True)
    output_file.write_text(f"SIGNAL:{signum}\\n")
    sys.exit(128 + signum)

signal.signal(signal.SIGINT, handler)
signal.signal(signal.SIGTERM, handler)

# Write ready marker
ready_file = Path(str(output_file) + ".ready")
ready_file.write_text(str(os.getpid()))
print(f"Child ready, pid={os.getpid()}", flush=True)

while True:
    print(".", end="", flush=True)
    time.sleep(0.2)
"""


def create_test_project(tmpdir: Path) -> Path:
    """Create a minimal pixi project for testing."""
    pixi_toml = tmpdir / "pixi.toml"
    pixi_toml.write_text("""
[project]
name = "signal-test"
channels = ["conda-forge"]
platforms = ["linux-64", "osx-arm64", "osx-64", "win-64"]

[tasks]
test = "python child.py"

[dependencies]
python = ">=3.10"
""")

    child_py = tmpdir / "child.py"
    child_py.write_text(CHILD_SCRIPT)

    return tmpdir


def run_test(name: str, test_func) -> bool:
    """Run a single test and report results."""
    print(f"\n{'='*60}")
    print(f"TEST: {name}")
    print('='*60)

    try:
        result = test_func()
        if result:
            print(f"{GREEN}✓ PASSED{RESET}")
            return True
        else:
            print(f"{RED}✗ FAILED{RESET}")
            return False
    except Exception as e:
        print(f"{RED}✗ FAILED with exception: {e}{RESET}")
        import traceback
        traceback.print_exc()
        return False


def wait_for_ready(ready_file: Path, timeout: float = 5.0) -> bool:
    """Wait for child process to signal it's ready."""
    start = time.time()
    while time.time() - start < timeout:
        if ready_file.exists():
            return True
        time.sleep(0.1)
    return False


def test_noninteractive_sigint(tmpdir: Path) -> bool:
    """
    Test: Non-interactive mode + kill -INT → should forward

    When stdin is not a TTY (piped), pixi should forward SIGINT to child.
    """
    output_file = tmpdir / "output.txt"
    ready_file = Path(str(output_file) + ".ready")
    output_file.unlink(missing_ok=True)
    ready_file.unlink(missing_ok=True)

    env = os.environ.copy()
    env["OUTPUT_FILE"] = str(output_file)

    # Run pixi with piped stdin (non-interactive)
    proc = subprocess.Popen(
        [str(PIXI_BIN), "run", "--manifest-path", str(tmpdir / "pixi.toml"), "test"],
        stdin=subprocess.PIPE,
        stdout=subprocess.PIPE,
        stderr=subprocess.PIPE,
        cwd=tmpdir,
        env=env,
    )

    # Wait for child to be ready
    if not wait_for_ready(ready_file):
        print(f"  {RED}Child did not start in time{RESET}")
        proc.kill()
        return False

    print(f"  Child is ready")

    # Send SIGINT to pixi
    print(f"  Sending SIGINT to pixi (pid={proc.pid})")
    os.kill(proc.pid, signal.SIGINT)

    # Wait for process to exit
    try:
        proc.wait(timeout=3)
    except subprocess.TimeoutExpired:
        proc.kill()
        print(f"  {RED}Process did not exit after SIGINT{RESET}")
        return False

    # Check if child received the signal
    time.sleep(0.2)
    if output_file.exists():
        content = output_file.read_text().strip()
        print(f"  Child output: {content}")
        if "SIGNAL:2" in content:  # SIGINT = 2
            print(f"  {GREEN}Child received SIGINT as expected{RESET}")
            return True

    print(f"  {RED}Child did not receive SIGINT{RESET}")
    return False


def test_noninteractive_sigterm(tmpdir: Path) -> bool:
    """
    Test: Non-interactive mode + kill -TERM → should forward

    SIGTERM should always be forwarded regardless of mode.
    """
    output_file = tmpdir / "output.txt"
    ready_file = Path(str(output_file) + ".ready")
    output_file.unlink(missing_ok=True)
    ready_file.unlink(missing_ok=True)

    env = os.environ.copy()
    env["OUTPUT_FILE"] = str(output_file)

    proc = subprocess.Popen(
        [str(PIXI_BIN), "run", "--manifest-path", str(tmpdir / "pixi.toml"), "test"],
        stdin=subprocess.PIPE,
        stdout=subprocess.PIPE,
        stderr=subprocess.PIPE,
        cwd=tmpdir,
        env=env,
    )

    if not wait_for_ready(ready_file):
        print(f"  {RED}Child did not start in time{RESET}")
        proc.kill()
        return False

    print(f"  Child is ready")
    print(f"  Sending SIGTERM to pixi (pid={proc.pid})")
    os.kill(proc.pid, signal.SIGTERM)

    try:
        proc.wait(timeout=3)
    except subprocess.TimeoutExpired:
        proc.kill()
        print(f"  {RED}Process did not exit after SIGTERM{RESET}")
        return False

    time.sleep(0.2)
    if output_file.exists():
        content = output_file.read_text().strip()
        print(f"  Child output: {content}")
        if "SIGNAL:15" in content:  # SIGTERM = 15
            print(f"  {GREEN}Child received SIGTERM as expected{RESET}")
            return True

    print(f"  {RED}Child did not receive SIGTERM{RESET}")
    return False


def test_interactive_sigterm(tmpdir: Path) -> bool:
    """
    Test: Interactive mode + kill -TERM → should forward

    SIGTERM should always be forwarded, even in interactive mode.
    We simulate interactive mode by using a PTY.
    """
    output_file = tmpdir / "output.txt"
    ready_file = Path(str(output_file) + ".ready")
    output_file.unlink(missing_ok=True)
    ready_file.unlink(missing_ok=True)

    env = os.environ.copy()
    env["OUTPUT_FILE"] = str(output_file)

    # Use PTY to simulate interactive mode (makes stdin a TTY)
    import pty

    master_fd, slave_fd = pty.openpty()

    proc = subprocess.Popen(
        [str(PIXI_BIN), "run", "--manifest-path", str(tmpdir / "pixi.toml"), "test"],
        stdin=slave_fd,
        stdout=subprocess.PIPE,
        stderr=subprocess.PIPE,
        cwd=tmpdir,
        env=env,
    )

    os.close(slave_fd)

    if not wait_for_ready(ready_file):
        print(f"  {RED}Child did not start in time{RESET}")
        proc.kill()
        os.close(master_fd)
        return False

    print(f"  Child is ready")
    print(f"  Sending SIGTERM to pixi (pid={proc.pid}) in interactive mode")
    os.kill(proc.pid, signal.SIGTERM)

    try:
        proc.wait(timeout=3)
    except subprocess.TimeoutExpired:
        proc.kill()
        print(f"  {RED}Process did not exit after SIGTERM{RESET}")
        os.close(master_fd)
        return False

    os.close(master_fd)

    time.sleep(0.2)
    if output_file.exists():
        content = output_file.read_text().strip()
        print(f"  Child output: {content}")
        if "SIGNAL:15" in content:
            print(f"  {GREEN}Child received SIGTERM as expected{RESET}")
            return True

    print(f"  {RED}Child did not receive SIGTERM{RESET}")
    return False


def test_interactive_sigint_same_pgid(tmpdir: Path) -> bool:
    """
    Test: Interactive mode + kill -INT + same PGID → should NOT forward

    This is the trade-off we accept (same as UV). When in interactive mode
    with same PGID, we assume CTRL+C was used and the terminal already
    delivered the signal to the child.

    Expected behavior: Child does NOT receive SIGINT from pixi's forwarding.
    (In real usage with CTRL+C, the child would get it from the terminal.)
    """
    output_file = tmpdir / "output.txt"
    ready_file = Path(str(output_file) + ".ready")
    output_file.unlink(missing_ok=True)
    ready_file.unlink(missing_ok=True)

    env = os.environ.copy()
    env["OUTPUT_FILE"] = str(output_file)

    import pty

    master_fd, slave_fd = pty.openpty()

    proc = subprocess.Popen(
        [str(PIXI_BIN), "run", "--manifest-path", str(tmpdir / "pixi.toml"), "test"],
        stdin=slave_fd,
        stdout=subprocess.PIPE,
        stderr=subprocess.PIPE,
        cwd=tmpdir,
        env=env,
    )

    os.close(slave_fd)

    if not wait_for_ready(ready_file):
        print(f"  {RED}Child did not start in time{RESET}")
        proc.kill()
        os.close(master_fd)
        return False

    print(f"  Child is ready")
    print(f"  Sending SIGINT to pixi (pid={proc.pid}) in interactive mode")
    print(f"  (Expected: child should NOT receive it - this is the UV trade-off)")
    os.kill(proc.pid, signal.SIGINT)

    # Give some time for signal to propagate (or not)
    time.sleep(0.5)

    # Check if child received the signal
    if output_file.exists():
        content = output_file.read_text().strip()
        print(f"  Child output: {content}")
        if "SIGNAL:2" in content:
            print(f"  {YELLOW}Child received SIGINT - this means forwarding happened{RESET}")
            print(f"  {YELLOW}(This would cause double-delivery with CTRL+C){RESET}")
            # This is actually a failure of the UV approach - we're forwarding when we shouldn't
            try:
                proc.wait(timeout=2)
            except:
                proc.kill()
            os.close(master_fd)
            return False

    # Child should still be running (didn't get our signal)
    if proc.poll() is None:
        print(f"  {GREEN}Child did NOT receive SIGINT (correct UV behavior){RESET}")
        print(f"  {GREEN}In real CTRL+C scenario, terminal would deliver it{RESET}")
        proc.kill()
        proc.wait()
        os.close(master_fd)
        return True

    # Process exited for some other reason
    os.close(master_fd)
    print(f"  {RED}Unexpected process exit{RESET}")
    return False


def main():
    print(f"Using pixi binary: {PIXI_BIN}")

    with tempfile.TemporaryDirectory() as tmpdir:
        tmpdir = Path(tmpdir)
        create_test_project(tmpdir)

        # Install dependencies first
        print(f"\n{YELLOW}Installing pixi environment...{RESET}")
        result = subprocess.run(
            [str(PIXI_BIN), "install", "--manifest-path", str(tmpdir / "pixi.toml")],
            cwd=tmpdir,
            capture_output=True,
            text=True,
        )
        if result.returncode != 0:
            print(f"{RED}Failed to install pixi environment:{RESET}")
            print(result.stderr)
            sys.exit(1)
        print(f"{GREEN}Environment ready{RESET}")

        # Run all tests
        results = []

        results.append(run_test(
            "Non-interactive + SIGINT → should forward",
            lambda: test_noninteractive_sigint(tmpdir)
        ))

        results.append(run_test(
            "Non-interactive + SIGTERM → should forward",
            lambda: test_noninteractive_sigterm(tmpdir)
        ))

        results.append(run_test(
            "Interactive + SIGTERM → should forward",
            lambda: test_interactive_sigterm(tmpdir)
        ))

        results.append(run_test(
            "Interactive + SIGINT + same PGID → should NOT forward (UV trade-off)",
            lambda: test_interactive_sigint_same_pgid(tmpdir)
        ))

        # Summary
        print(f"\n{'='*60}")
        print("SUMMARY")
        print('='*60)
        passed = sum(results)
        total = len(results)

        if passed == total:
            print(f"{GREEN}All {total} tests passed!{RESET}")
        else:
            print(f"{RED}{total - passed} of {total} tests failed{RESET}")

        sys.exit(0 if passed == total else 1)


if __name__ == "__main__":
    main()